package cn.itcast.tags.models.ml

import cn.itcast.tags.models.{AbstractModel, ModelType}
import cn.itcast.tags.tools.TagTools
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.ml.clustering.{KMeans, KMeansModel}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.linalg
import org.apache.spark.sql.{Column, DataFrame, SparkSession}
import org.apache.spark.sql.functions.{when, _}
import org.apache.spark.sql.types.{DataTypes, DoubleType}
import org.apache.spark.storage.StorageLevel

/**
 * 挖掘类型标签：客户价值模型RFM
 *    R:消费周期
 *    F:消费次数
 *    M:消费金额
 */
class RfmModel extends AbstractModel("客户价值标签",ModelType.ML){
  /*
   高价值0---------低价值4
   */
  override def doTag(businessDF: DataFrame, tagDF: DataFrame): DataFrame = {
    val session: SparkSession = businessDF.sparkSession
    import session.implicits._
    //TODO 1.计算每个用户RFM值，按照用户memberid分组，再进行聚合函数聚合统计
    //R:消费周期,finishtime,日期时间函数current_timestamp,from_unixtimestamp,datediff
    //F:消费次数 ordersn 函数count
    //M:消费金额 orderamount 函数sum
    
    val rfmDF: DataFrame = businessDF
      .groupBy($"memberid")
      .agg(
        max($"finishtime").as("max_finishtime"),
        count("ordersn").as("frequency"),
        sum(
          $"orderamount".cast(DataTypes.createDecimalType(10,2))
        ).as("monetary")
      )
      .select(
        $"memberid".as("userId"),
        //计算R值
        datediff(
          current_timestamp(),from_unixtime($"max_finishtime")
        ).as("recency"),
        $"frequency",
        $"monetary"
      )

    rfmDF.show(10,false)

    //TODO 2.按照规则该RFM进行打分（RFM_SCORD)
    //R: 1-3天=5分，4-6天=4,7-9天=3分，10-15天=2分，大于16天=1分
    //F: >=200=5分，150-199=4分，100-149=3分，50-99=2分，1-49=1分
    //M: >=20w=5分，10-19w=4分，5-9w=3分，1-4w=2分，<1w=1分
    //使用CASE WHEN...WHEN...ELSE...END

    val rColumn: Column =when($"recency".between(1,3),5.0)
      .when($"recency".between(4,6),4.0)
      .when($"recency".between(7,9),3.0)
      .when($"recency".between(10,15),2.0)
      .when($"recency".geq(16),1.0)


    val fColumn: Column =when($"frequency".between(1,49),1.0)
      .when($"frequency".between(50,99),2.0)
      .when($"frequency".between(100,149),3.0)
      .when($"frequency".between(150,199),4.0)
      .when($"frequency".geq(200),5.0)

    val mColumn: Column =when($"monetary".geq(200000),5.0)
      .when($"monetary".between(100000,199999),4.0)
      .when($"monetary".between(50000,99999),3.0)
      .when($"monetary".between(10000,49999),2.0)
      .when($"monetary".lt(10000),1.0)

    val rfmScoreDF: DataFrame = rfmDF.select(
      $"userId",
      rColumn.as("r_score"),
      fColumn.as("f_score"),
      mColumn.as("m_score")
    )


    // 4. 组合R\F\M列为特征值features
    val assembler: VectorAssembler = new VectorAssembler()
      .setInputCols(Array("r_score", "f_score", "m_score"))
      .setOutputCol("features")
    val featureDF: DataFrame = assembler.transform(rfmScoreDF)

    // 将训练数据缓存
    featureDF.persist(StorageLevel.MEMORY_AND_DISK)

    //3.2使用k-meams算法进行聚类，训练模型
    val kMeansModel: KMeansModel = trainModel(featureDF)


    //模型预测
    val predictionDF: DataFrame = kMeansModel.transform(featureDF)

    predictionDF.show(50,false)

    val indexTagMap: Map[Int, String] = TagTools.convertIndexMap(kMeansModel.clusterCenters, tagDF)

    println("==============")
    indexTagMap.foreach(println) //(列簇索引,rfm值)

    //使用预测值打标签
    val indexTagMapBroadcast: Broadcast[Map[Int, String]] = session.sparkContext.broadcast(indexTagMap)
    val index_to_tag = udf(
      (clusterIndex:Int)=>indexTagMapBroadcast.value(clusterIndex)
    )
    val modelDF: DataFrame = predictionDF.select(
      $"userId",
      index_to_tag($"prediction").as("rfm")
    )
    //modelDF.show(100,false)
    null
  }

  def trainModel(dataframe:DataFrame): KMeansModel ={
    val kMeansModel: KMeansModel = new KMeans()
      .setFeaturesCol("features")
      .setPredictionCol("prediction") //由于k=5,所以预测值prediction范围：0,1,2,3,4
      .setK(5)
      .setMaxIter(10)
      .fit(dataframe)
    println(s"WSSSE=${kMeansModel.computeCost(dataframe)}")
    kMeansModel
  }
}


object RfmModel{
  def main(args: Array[String]): Unit = {
    val model = new RfmModel
    model.executeModel(361L,false)
  }
}